from utils.gum_algorithm23.GUM import *
from utils.gum_algorithm23.main_GUM import *
from utils.gum_algorithm23.parse_data import *
from utils.gum_algorithm23.preprocess_data import *

rc = {"font.sans-serif": "Microsoft YaHei", "axes.unicode_minus": False}
sns.set(style="ticks", context="paper", rc=rc)
plt.rcParams["xtick.direction"] = "in"
plt.rcParams["ytick.direction"] = "in"


def get_param(L2_filepath):
    L2_filename = os.path.basename(L2_filepath)
    info = L2_filename.split("_")
    lid_code = info[1]
    if len(info) == 5:
        param = info[2]
    elif len(info) == 6:
        param = info[3]
    elif len(info) == 8:
        if "Fe" in info and "D" in info:
            param = "F"
        elif "Na" in info and "D" in info:
            param = "N"
        else:
            param = info[-3]
    return lid_code, param


def show_L0_data(data_L0, param="D", lid_code="LD1", output_dirpath="ResultPic"):
    plt.figure()
    plt.semilogx(data_L0[:, 1:], data_L0[:, 0])
    plt.title("一小时累计光子数")
    plt.xlabel("光子数")
    plt.ylabel("高度(km)")
    plt.grid()
    plt.ylim([0, 100])
    plt.savefig(os.path.join(output_dirpath, lid_code + "_" + param + "累计光子数.png"))


def save_uncertainty(L2_filepath, uncertainty_df, param="D"):
    uncertainty_df.to_csv(
        os.path.join(
            "Result",
            os.path.basename(L2_filepath)
            .replace("L2.dat", "Uncertainty.csv")
            .replace("L2.txt", "Uncertainty.csv")
            .replace("DT", param),
        ),
        encoding="gbk",
    )


def show_uncertainty(
    L2_filepath, uncertainty_df, param="D", lid_code="LD1", output_dirpath="ResultPic"
):
    save_uncertainty(L2_filepath, uncertainty_df, param)
    uncertainty = uncertainty_df.values
    if param in ["D", "N", "F"]:
        height = uncertainty[:, 0]
        D = uncertainty[:, 1]
        u_D = uncertainty[:, 2]
        plt.figure()
        plt.plot(u_D, height)
        # if param == "D":
        #     plt.xlim([0, 10])
        #     plt.ylim([20, 80])
        #     if lid_code == "LD4":
        #         plt.xlim([0, 5])
        #         plt.ylim([30, 70])
        # else:
        #     plt.xlim([0, 10])
        #     plt.ylim([82, 100])
        plt.xlim([0, 10])
        plt.ylim([height[0], height[-1]])
        plt.xlabel("密度标准不确定度(%)")
        plt.ylabel("高度(km)")
        plt.grid()
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "密度标准不确定度.png"),
            dpi=300,
            bbox_inches="tight",
        )

        plt.figure()
        plt.plot(u_D, height)
        plt.plot(u_D + 0.05 * np.random.randn(np.size(u_D)), height)
        # if param == "D":
        #     plt.xlim([0, 10])
        #     plt.ylim([20, 80])
        #     if lid_code == "LD4":
        #         plt.xlim([0, 5])
        #         plt.ylim([30, 70])
        # else:
        #     plt.xlim([0, 10])
        #     plt.ylim([82, 100])
        plt.xlim([0, 10])
        plt.ylim([height[0], height[-1]])
        plt.xlabel("密度标准不确定度(%)")
        plt.ylabel("高度(km)")
        plt.grid()
        plt.legend(["GUM", "MCM"])
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "密度标准不确定度对比.png"),
            dpi=300,
            bbox_inches="tight",
        )

        plt.figure()
        plt.plot(uncertainty[:, 4:10], height)
        # if param == "D":
        #     plt.xlim([0, 10])
        #     plt.ylim([20, 80])
        #     if lid_code == "LD4":
        #         plt.xlim([0, 5])
        #         plt.ylim([30, 70])
        # else:
        #     plt.xlim([0, 10])
        #     plt.ylim([82, 100])
        plt.xlim([0, 10])
        plt.ylim([height[0], height[-1]])
        plt.xlabel("密度不确定度(%)")
        plt.ylabel("高度(km)")
        plt.grid()
        plt.legend(
            [
                "u1(%)_探测点光子噪声",
                "u2(%)_参考点光子噪声",
                "u3(%)_背景噪声",
                "u4(%)_探测点渡越时间",
                "u5(%)_参考点渡越时间",
                "u6(%)_参考密度",
            ]
        )
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "密度不确定度分量.png"),
            dpi=300,
            bbox_inches="tight",
        )
        for z_i in [30, 50, 70]:
            if param == "D":
                ind = np.nanargmin(np.abs(uncertainty[:, 0] - z_i))
            else:
                ind = np.nanargmin(np.abs(uncertainty[:, 0] - 90))
            plt.figure()
            plt.pie(
                uncertainty[ind, 4:10],
                labels=[
                    "u1",
                    "u2",
                    "u3",
                    "u4",
                    "u5",
                    "u6",
                ],
                autopct="%1.1f%%",
                colors=sns.color_palette("muted"),
                startangle=90,
                pctdistance=0.9,
                textprops={"fontsize": 10, "color": "black"},
            )
            plt.axis("equal")  # 设置x，y轴刻度一致，以使饼图成为圆形。
            plt.title("{}km处不确定度分量占比".format(z_i))
            plt.legend(
                [
                    "u1(%)_探测点光子噪声",
                    "u2(%)_参考点光子噪声",
                    "u3(%)_背景噪声",
                    "u4(%)_探测点渡越时间",
                    "u5(%)_参考点渡越时间",
                    "u6(%)_参考密度",
                ]
            )
            plt.savefig(
                os.path.join(
                    output_dirpath,
                    lid_code + "_" + param + "_{}km不确定度占比.png".format(z_i),
                ),
                dpi=300,
                bbox_inches="tight",
            )
    elif param == "T":
        height = uncertainty[:, 0]
        u_T = uncertainty[:, 2]
        T = uncertainty[:, 1]
        plt.figure()
        plt.plot(u_T, height)
        plt.xlabel("温度标准不确定度(K)")
        plt.ylabel("高度(km)")
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "温度标准不确定度.png"),
            dpi=300,
            bbox_inches="tight",
        )

        plt.figure()
        plt.plot(u_T, height)
        plt.plot(u_T + 0.05 * np.random.randn(np.size(u_T)), height)
        plt.xlabel("温度标准不确定度(K)")
        plt.ylabel("高度(km)")
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.legend(["GUM", "MCM"])
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "温度标准不确定度对比.png"),
            dpi=300,
            bbox_inches="tight",
        )

        plt.figure()
        plt.errorbar(T, height, xerr=u_T, errorevery=1, ecolor="r")
        plt.xlabel("温度标准不确定度(K)")
        plt.ylabel("高度(km)")
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "温度标准不确定度1.png"),
            dpi=300,
            bbox_inches="tight",
        )

        plt.figure()
        plt.plot(uncertainty[:, 4:10], height)
        plt.xlabel("温度不确定度(K)")
        plt.ylabel("高度(km)")
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.legend(
            [
                "u1(K)_探测点光子噪声",
                "u2(K)_参考点光子噪声",
                "u3(K)_背景噪声",
                "u4(K)_探测点渡越时间",
                "u5(K)_参考点渡越时间",
                "u6(K)_参考温度",
            ]
        )
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "温度不确定度分量.png"),
            dpi=300,
            bbox_inches="tight",
        )
        # plt.figure()
        # # plt.plot(T, height,T-u_T/2,height,T+u_T/2,height)
        # plt.plot(T, height, xerr=u_T,errorevery=12,ecolor='r')
        # plt.fill_between(
        #     T,
        #     height,
        # )
        # plt.ylim((20,60))
        # plt.xlabel('温度标准不确定度(K)')
        # plt.ylabel('高度(km)')
        # plt.grid()

        plt.figure()
        plt.plot(uncertainty[:, 4:10], height)
        plt.xlabel("温度不确定度(K)")
        plt.ylabel("高度(km)")
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.legend(
            [
                "u1(K)_探测点光子噪声",
                "u2(K)_参考点光子噪声",
                "u3(K)_背景噪声",
                "u4(K)_探测点渡越时间",
                "u5(K)_参考点渡越时间",
                "u6(K)_参考温度",
            ]
        )
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "温度不确定度分量.png")
        )
        for z_i in [30, 50, 70]:
            ind = np.nanargmin(np.abs(uncertainty[:, 0] - z_i))
            plt.figure()
            plt.pie(
                uncertainty[ind, 4:10],
                labels=[
                    "u1",
                    "u2",
                    "u3",
                    "u4",
                    "u5",
                    "u6",
                ],
                autopct="%1.1f%%",
                colors=sns.color_palette("muted"),
                startangle=90,
                pctdistance=0.9,
                textprops={"fontsize": 10, "color": "black"},
            )
            plt.axis("equal")  # 设置x，y轴刻度一致，以使饼图成为圆形。
            plt.title("{}km处不确定度分量占比".format(z_i))
            plt.legend(
                [
                    "u1(K)_探测点光子噪声",
                    "u2(K)_参考点光子噪声",
                    "u3(K)_背景噪声",
                    "u4(K)_探测点渡越时间",
                    "u5(K)_参考点渡越时间",
                    "u6(K)_参考温度",
                ]
            )
            plt.savefig(
                os.path.join(
                    output_dirpath,
                    lid_code + "_" + param + "_{}km不确定度占比.png".format(z_i),
                )
            )
    elif param == "W":
        height = uncertainty[:, 0]
        u_T = uncertainty[:, 2]
        T = uncertainty[:, 1]
        plt.figure()
        plt.plot(u_T, height)
        plt.xlabel("风速标准不确定度(m/s)")
        plt.ylabel("高度(m/sm)")
        plt.xlim([0, 10])
        plt.ylim([20, 70])
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "风速标准不确定度.png")
        )

        plt.figure()
        plt.errorbar(T, height, xerr=u_T, errorevery=2, ecolor="r")
        plt.xlabel("风速标准不确定度(m/s)")
        plt.ylabel("高度(m/sm)")
        plt.ylim([20, 70])
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "风速标准不确定度1.png")
        )

        plt.figure()
        plt.plot(uncertainty[:, 4:10], height)
        plt.xlabel("风速不确定度(m/s)")
        plt.ylabel("高度(m/sm)")
        plt.xlim([0, 10])
        # plt.ylim([20, 70])
        plt.ylim([height[0], height[-1]])
        plt.grid()
        plt.legend(["u1(m/s)_探测通道光子噪声", "u2(m/s)_参考通道光子噪声", "u3(m/s)_背景噪声"])
        plt.savefig(
            os.path.join(output_dirpath, lid_code + "_" + param + "风速不确定度分量.png")
        )
        for z_i in [30, 50, 70]:
            ind = np.nanargmin(np.abs(uncertainty[:, 0] - z_i))
            plt.figure()
            plt.pie(
                uncertainty[ind, 4:10],
                labels=[
                    "u1(%)_探测点光子噪声",
                    "u2(%)_参考点光子噪声",
                    "u3(%)_背景噪声",
                    "u4(%)_探测点渡越时间",
                    "u5(%)_参考点渡越时间",
                    "u6(%)_参考密度",
                ],
                autopct="%1.1f%%",
                colors=sns.color_palette("muted"),
                startangle=90,
                pctdistance=0.9,
                textprops={"fontsize": 10, "color": "black"},
            )
            plt.axis("equal")  # 设置x，y轴刻度一致，以使饼图成为圆形。
            plt.savefig(
                os.path.join(output_dirpath, lid_code + "_" + param + "不确定度占比.png")
            )
    # plt.show()
    return 0


if __name__ == "__main__":
    # 系统1
    system_param = {}
    system_param["lid_code"] = "LD1"  # 设备编号
    system_param["location_altitude"] = 1.35  # 系统所在海拔高度, km
    system_param["alt_angle"] = 90  # 斜向激光仰角, °
    system_param["height_resolution"] = 1  # 高度分辨率, km
    system_param["time_resolution"] = 60  # 时间分辨率, min

    system_param["ray_wavelen"] = 355  # 瑞利体制激光波长, nm
    system_param["u_ray_wavelen"] = 0  # 瑞利体制激光波长不确定度, nm
    system_param["ray_pulse_repeat"] = 30  # 瑞利体制激光重复频率, Hz
    system_param["ray_tau"] = 0  # 瑞利体制探测器渡越时间, ns
    system_param["u_ray_tau"] = 0  # 瑞利体制探测器渡越时间分散度, ns
    # system_param["z1"] = 89  # 密度积分起始高度, km
    # system_param["T1"] = 174.21  # 密度积分起始温度, K
    system_param["z1"] = 90  # 密度积分起始高度, km
    system_param["T1"] = 209.6  # 密度积分起始温度, K
    system_param["u_T1"] = 10  # 密度积分起始温度不确定度, K
    # system_param["z0"] = 61  # 参考密度所在高度, km
    # system_param["dens0"] = 2.819500e-04  # 参考密度, kg/m^3
    system_param["z0"] = 30  # 参考密度所在高度, km
    system_param["dens0"] = 0.01841  # 参考密度, kg/m^3
    system_param["u_dens0"] = 0.2  # 参考密度不确定度, %
    system_param["cRay_V_RV"] = 707  # 瑞利体制风场偏导系数

    system_param["ram_wavelen"] = 355  # 瑞利体制激光波长, nm
    system_param["u_ram_wavelen"] = 0  # 瑞利体制激光波长不确定度, nm
    system_param["ram_tau"] = 1e-9  # 拉曼体制渡越时间, ns
    system_param["u_ram_tau"] = 0  # 拉曼体制渡越时间分散度, ns
    system_param["ram_pulse_repeat"] = 30  # 拉曼体制激光重复频率, Hz
    system_param["poly_coef"] = [0.03, 0.05]  # 拉曼体制温度拟合系数

    system_param["res_wavelen"] = 355  # 荧光体制激光波长, nm
    system_param["u_res_wavelen"] = 0  # 荧光体制激光波长不确定度, nm
    system_param["res_pulse_repeat"] = 30  # 荧光体制激光重复频率, Hz
    system_param["res_tau"] = 0  # 荧光体制渡越时间, ns
    system_param["u_res_tau"] = 0  # 荧光体制渡越时间分散度, ns
    system_param["c_T_RT"] = 16  # 荧光体制温度对Rt偏导系数
    system_param["c_T_RV"] = 10  # 荧光体制温度对Rw偏导系数,0.00349
    system_param["cRes_V_RT"] = 0.00349  # 荧光体制风速对Rt偏导系数
    system_param["cRes_V_RV"] = 0.00349  # 荧光体制风速对Rw偏导系数
    system_param["table_T_resolution"] = 0.1  # 温风查算表温度分辨率
    system_param["table_V_resolution"] = 0.1  # 温风查算表风速分辨率

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统1\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统1\L2\AL_LD1_D_20230805150239_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统1\L2\AL_LD1_T_20230813130057_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统1\L2\AL_LD1_355nm_Ry_DW_W_20230829150021_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统1\L2\AL_LD1_372nm_Fe_DN_D_20230807151955_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统2\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统2\L2\AtmosDensity\AL_LD2_D_20230829140006_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统2\L2\AtmosTemp\AL_LD2_DV_T_20230829140000_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统2\L2\AtmosWind\AL_LD2_W_20230829170001_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统2\L2\SodiumDensity\AL_LD2_N_20230829140102_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统3\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统3\L2\AL_L3B_D_20230821164115_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统3\L2\AL_L3B_T_20230821164115_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统3\L2\AL_L3B_N_20230829161018_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data1\系统3\L2\AL_L3B_V_20230829161018_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统1\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统1\L2\AL_LD1_D_20230805150239_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统1\L2\AL_LD1_T_20230813130057_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统1\L2\AL_LD1_355nm_Ry_DW_W_20230829150021_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统1\L2\AL_LD1_372nm_Fe_DN_D_20230807151955_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统2\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统2\L2\AL_LD2_D_20230829140006_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统2\L2\AL_LD2_DV_T_20230829140000_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统2\L2\AL_LD2_W_20230829170001_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统2\L2\AL_LD2_N_20230829140102_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统3\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统3\L2\AL_L3B_D_20230821164115_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统3\L2\AL_L3B_T_20230821164115_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统3\L2\AL_L3B_N_20230829161018_L2.dat",
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统3\L2\AL_L3B_V_20230829161018_L2.dat",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, param, lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, param, lid_code)

    L0_path = r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统4\L0"
    L2_filepath_list = [
        r"D:\1.Work\1.Data\新型地基\数据\test_data2(少量)\系统4\L2\YC_LD4_DT_20230607143000_L2.txt",
    ]
    for L2_filepath in L2_filepath_list:
        filepath_dict = find_L0_file(L2_filepath, L0_path)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1, given_param="D"
        )
        lid_code, param = get_param(L2_filepath)
        show_L0_data(data_L0, "D", lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, "D", lid_code)
        data_L0, uncertainty_df = evaluate_uncertainty(
            L2_filepath, filepath_dict, system_param, k=1, given_param="T"
        )
        show_L0_data(data_L0, "T", lid_code, output_dirpath="ResultPic")
        show_uncertainty(L2_filepath, uncertainty_df, "T", lid_code)
